package com.zys.background.auth.config;

import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.source.ImmutableJWKSet;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;
import com.zys.background.auth.extension.password.OAuth2PasswordAuthenticationConverter;
import com.zys.background.auth.extension.password.OAuth2PasswordAuthenticationProvider;
import com.zys.background.auth.handler.OAuth2LogoutHandler;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.config.Customizer;
import org.springframework.security.config.annotation.authentication.builders.AuthenticationManagerBuilder;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.annotation.web.configurers.oauth2.server.resource.OAuth2ResourceServerConfigurer;
import org.springframework.security.config.http.SessionCreationPolicy;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.userdetails.User;
import org.springframework.security.crypto.bcrypt.BCryptPasswordEncoder;
import org.springframework.security.crypto.password.PasswordEncoder;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.server.authorization.JdbcOAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.client.JdbcRegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configuration.OAuth2AuthorizationServerConfiguration;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2AuthorizationServerConfigurer;
import org.springframework.security.oauth2.server.authorization.settings.AuthorizationServerSettings;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.security.oauth2.server.authorization.web.authentication.DelegatingAuthenticationConverter;
import org.springframework.security.oauth2.server.authorization.web.authentication.OAuth2AuthorizationCodeAuthenticationConverter;
import org.springframework.security.oauth2.server.authorization.web.authentication.OAuth2ClientCredentialsAuthenticationConverter;
import org.springframework.security.oauth2.server.authorization.web.authentication.OAuth2RefreshTokenAuthenticationConverter;
import org.springframework.security.web.AuthenticationEntryPoint;
import org.springframework.security.web.SecurityFilterChain;
import org.springframework.security.web.authentication.AuthenticationFailureHandler;
import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
import org.springframework.security.web.authentication.logout.LogoutSuccessHandler;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.*;

/**
 * @author zys
 * @since 2022-11-21
 */
@Configuration
@RequiredArgsConstructor
public class AuthorizationServerConfiguration {
  private final JdbcTemplate jdbcTemplate;
  private final LogoutSuccessHandler logoutSuccessHandler;
  private final AuthenticationEntryPoint authenticationEntryPoint;
  private final AuthenticationFailureHandler authenticationFailureHandler;
  private final AuthenticationSuccessHandler authenticationSuccessHandler;

  @Bean
  @Order(Ordered.HIGHEST_PRECEDENCE)
  public SecurityFilterChain authorizationServerSecurityFilterChain(
      HttpSecurity http,
      AuthenticationManager authenticationManager,
      RegisteredClientRepository registeredClientRepository,
      OAuth2AuthorizationService oAuth2AuthorizationService
  ) throws Exception {
    OAuth2AuthorizationServerConfiguration.applyDefaultSecurity(http);

    DelegatingAuthenticationConverter converter = new DelegatingAuthenticationConverter(
        Arrays.asList(
            new OAuth2AuthorizationCodeAuthenticationConverter(),
            new OAuth2RefreshTokenAuthenticationConverter(),
            new OAuth2ClientCredentialsAuthenticationConverter(),
            new OAuth2PasswordAuthenticationConverter()
        )
    );


    http.getConfigurer(OAuth2AuthorizationServerConfigurer.class)
        // 默认配置的认证路径
        .authorizationServerSettings(AuthorizationServerSettings.builder().build())
        // 配置 token 生成逻辑 .tokenGenerator()
        // 配置认证逻辑 .authorizationService()
        .authorizationService(oAuth2AuthorizationService)
        // 配置授权逻辑 .authorizationConsentService()
        //.authorizationConsentService(new JdbcOAuth2AuthorizationConsentService(jdbcTemplate,registeredClientRepository))
        // 客户端身份认证 .clientAuthentication()
        // 配置已注册客户端处理
        .tokenEndpoint(tokenEndpoint -> tokenEndpoint
            .accessTokenRequestConverter(converter)
            .accessTokenResponseHandler(authenticationSuccessHandler)
            .errorResponseHandler(authenticationFailureHandler)
        )
        .registeredClientRepository(registeredClientRepository)
        .oidc(Customizer.withDefaults())
        .and()
        .exceptionHandling().authenticationEntryPoint(authenticationEntryPoint)
        .and()
        .oauth2ResourceServer(OAuth2ResourceServerConfigurer::jwt).build();

    OAuth2TokenGenerator<?> tokenGenerator = http.getSharedObject(OAuth2TokenGenerator.class);
    OAuth2PasswordAuthenticationProvider provider = new OAuth2PasswordAuthenticationProvider(authenticationManager, oAuth2AuthorizationService, tokenGenerator);

    return http.authenticationProvider(provider).getOrBuild();
  }


  @Bean
  @Order(2)
  public SecurityFilterChain defaultSecurityFilterChain(HttpSecurity http,OAuth2AuthorizationService oAuth2AuthorizationService) throws Exception {
    http.csrf().disable()
        .cors().disable()
        .sessionManagement().sessionCreationPolicy(SessionCreationPolicy.STATELESS)
        .and()
        .authorizeHttpRequests().anyRequest().authenticated()
        .and()
        .formLogin(Customizer.withDefaults())
        .logout().addLogoutHandler(new OAuth2LogoutHandler(oAuth2AuthorizationService)).logoutSuccessHandler(logoutSuccessHandler);
    return http.build();
  }

  @Bean
  public AuthenticationManager authenticationManager(AuthenticationManagerBuilder authenticationManagerBuilder) {
    return authenticationManagerBuilder.getOrBuild();
  }

  @Bean
  public RegisteredClientRepository registeredClientRepository() {
    return new JdbcRegisteredClientRepository(jdbcTemplate);
  }

  @Bean
  public OAuth2AuthorizationService oAuth2AuthorizationService(RegisteredClientRepository registeredClientRepository) {
    return new JdbcOAuth2AuthorizationService(jdbcTemplate, registeredClientRepository);
  }

  @Bean
  public PasswordEncoder passwordEncoder() {
    return new BCryptPasswordEncoder();
  }

  @Bean
  public OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer() {
    return context -> {
      // 客户端模式不返回具体用户信息
      if (!AuthorizationGrantType.CLIENT_CREDENTIALS.equals(context.getAuthorizationGrantType())) {
        JwtClaimsSet.Builder builder = context.getClaims();
        User user = (User) context.getPrincipal().getPrincipal();
        Collection<GrantedAuthority> authorities = user.getAuthorities();
        Set<String> res = new HashSet<>();
        for (GrantedAuthority authority : authorities) {
          res.add(authority.getAuthority());
        }
        builder.claims(claims -> {
          claims.put("username", user.getUsername());
          claims.put("authorities", res);
        });
      }
    };
  }

  @Bean
  public JWKSource<SecurityContext> jwkSource() {
    KeyPair keyPair = generateRsaKey();
    RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
    RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
    RSAKey rsaKey = new RSAKey.Builder(publicKey)
        .privateKey(privateKey)
        .keyID(UUID.randomUUID().toString())
        .build();
    JWKSet jwkSet = new JWKSet(rsaKey);
    return new ImmutableJWKSet<>(jwkSet);
  }

  @Bean
  public JwtDecoder jwtDecoder(JWKSource<SecurityContext> jwkSource) {
    return OAuth2AuthorizationServerConfiguration.jwtDecoder(jwkSource);
  }

  private static KeyPair generateRsaKey() {
    KeyPair keyPair;
    try {
      KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
      keyPairGenerator.initialize(2048);
      keyPair = keyPairGenerator.generateKeyPair();
    } catch (Exception ex) {
      throw new IllegalStateException(ex);
    }
    return keyPair;
  }
}
